from PIL import Image, ImageSequence, ImageOps
import torch
import os
import numpy as np
import urllib3
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from io import BytesIO
import folder_paths
import node_helpers

# 禁用 SSL 验证警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
# 默认请求头，模拟浏览器访问，防止部分服务器拒绝
DEFAULT_HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
}


class LoadImageUrlLocalCache:
    @classmethod
    def INPUT_TYPES(s):
        # input_dir = folder_paths.get_input_directory()
        # files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
        # files = folder_paths.filter_files_content_types(files, ["image"])
        # files = sorted(files)
        return {"required":
            {
                "image": ("COMBO", {"image_upload": True}),
                "url": ("STRING", {"default": "http://", "multiline": False})
            },
        }

    CATEGORY = "image"
    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load"
    DESCRIPTION = (
        "京科华安图片加载节点-V9 ,先从本地加载图片，如果不存在，再从url加载图片."
    )

    def load(self, image, url):
        image_path = folder_paths.get_annotated_filepath(image)
        img, fname = load_image(url, image_path)

        output_images = []
        output_masks = []
        w, h = None, None

        excluded_formats = ['MPO']

        for i in ImageSequence.Iterator(img):
            i = node_helpers.pillow(ImageOps.exif_transpose, i)

            if i.mode == 'I':
                i = i.point(lambda i: i * (1 / 255))
            image = i.convert("RGB")

            if len(output_images) == 0:
                w = image.size[0]
                h = image.size[1]

            if image.size[0] != w or image.size[1] != h:
                continue

            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
            if 'A' in i.getbands():
                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            elif i.mode == 'P' and 'transparency' in i.info:
                mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
            output_images.append(image)
            output_masks.append(mask.unsqueeze(0))

        if len(output_images) > 1 and img.format not in excluded_formats:
            output_image = torch.cat(output_images, dim=0)
            output_mask = torch.cat(output_masks, dim=0)
        else:
            output_image = output_images[0]
            output_mask = output_masks[0]

        return (output_image, output_mask)

    def check_lazy_status(self, image, url):
        image_path = folder_paths.get_annotated_filepath(image)
        needed = []
        if os.path.exists(image_path):
            print(f"存在文件：{image} -- {image_path}")
            needed.append("image")
        else:
            print(f"不不存在文件：{image} -- {image_path}")
        return needed


def load_image(url: str, image_path: str, retries: int = 3, timeout: float = 10.0):
    """
    image_path:本地图片路径
    支持本地路径或 URL 加载图片。
    - 本地路径优先
    - 网络加载支持域名/IP, SSL 验证跳过, 重试机制
    - 忽略环境代理，确保直接访问
    - 添加浏览器 User-Agent 头，防止 403 拒绝
    - 放宽 Content-Type 检查，允许 application/octet-stream
    返回 PIL Image 对象及文件名
    """
    # 本地文件
    if os.path.exists(image_path):
        print(f"加载换装图片使用本地路径：{image_path}")
        img = Image.open(image_path)
        filename = os.path.basename(image_path)
    else:
        print(f"加载换装图片使用url：{url}")
        input_dir = folder_paths.get_input_directory()
        source = url.strip()
        filename = os.path.basename(source.split('?', 1)[0]) or 'downloaded_image'
        # 网络加载
        session = requests.Session()
        session.trust_env = False
        session.headers.update(DEFAULT_HEADERS)
        retry_strategy = Retry(
            total=retries,
            backoff_factor=0.3,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["GET"]
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("https://", adapter)
        session.mount("http://", adapter)

        resp = session.get(source, timeout=timeout, verify=False)
        resp.raise_for_status()

        # 放宽 Content-Type 检查
        ctype = resp.headers.get('Content-Type', '').lower()
        # 部分服务器将图片标记为 application/octet-stream，也允许
        if ctype and not (ctype.startswith('image/') or 'octet-stream' in ctype):
            raise RuntimeError(f"URL does not point to an image, content-type: {ctype}")

        data = resp.content
        # Save image to input directory
        save_path = os.path.join(input_dir, filename)
        with open(save_path, "wb") as f:
            f.write(data)
        # Verify and load the saved image
        if not os.path.exists(save_path):
            raise FileNotFoundError(f"Failed to save image to {save_path}")
        try:
            img = Image.open(save_path)
        except Exception:
            # 尝试强制以 RGB 打开
            img = Image.open(save_path).convert("RGB")

    return img, filename
